/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "csr_matvec.h"

namespace YHAMG
{

void CSR_MatVec(double alpha, const CSR_Matrix& A, const Vector& x, double beta, const Vector& y, const Vector& w)
{
	int n = A.size[0];
	int* Ap = A.rowptr;
	int* Ai = A.colind;
	double* Av = A.values;
	double* xv = x.values;
	double* yv = y.values;
	double* wv = w.values;

	if (alpha == 0)
	{
		for (int i = 0; i < n; ++i)
			wv[i] = yv[i];
		if (beta != 1)
#ifdef USE_OPENMP
#pragma omp parallel for 
#endif
			for (int i = 0; i < n; ++i)
				wv[i] *= beta;
	}
	else if (alpha == 1)
	{
		if (beta == 0)
		{
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
			for (int i = 0; i < n; ++i)
			{
				double temp = 0.0;
				for (int j = Ap[i]; j < Ap[i + 1]; ++j)
					temp += Av[j] * xv[Ai[j]];
				wv[i] = temp;
			}
		}
		else if (beta == 1)
		{
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
			for (int i = 0; i < n; ++i)
			{
				double temp = yv[i];
				for (int j = Ap[i]; j < Ap[i + 1]; ++j)
					temp += Av[j] * xv[Ai[j]];
				wv[i] = temp;
			}
		}
		else if (beta == -1)
		{
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
			for (int i = 0; i < n; ++i)
			{
				double temp = -yv[i];
				for (int j = Ap[i]; j < Ap[i + 1]; ++j)
					temp += Av[j] * xv[Ai[j]];
				wv[i] = temp;
			}
		}
		else
		{
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
			for (int i = 0; i < n; ++i)
			{
				double temp = beta * yv[i];
				for (int j = Ap[i]; j < Ap[i + 1]; ++j)
					temp += Av[j] * xv[Ai[j]];
				wv[i] = temp;
			}
		}
	}
	else if (alpha == -1)
	{
		if (beta == 0)
		{
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
			for (int i = 0; i < n; ++i)
			{
				double temp = 0.0;
				for (int j = Ap[i]; j < Ap[i + 1]; ++j)
					temp -= Av[j] * xv[Ai[j]];
				wv[i] = temp;
			}
		}
		else if (beta == 1)
		{
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
			for (int i = 0; i < n; ++i)
			{
				double temp = yv[i];
				for (int j = Ap[i]; j < Ap[i + 1]; ++j)
					temp -= Av[j] * xv[Ai[j]];
				wv[i] = temp;
			}
		}
		else if (beta == -1)
		{
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
			for (int i = 0; i < n; ++i)
			{
				double temp = -yv[i];
				for (int j = Ap[i]; j < Ap[i + 1]; ++j)
					temp -= Av[j] * xv[Ai[j]];
				wv[i] = temp;
			}
		}
		else
		{
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
			for (int i = 0; i < n; ++i)
			{
				double temp = beta * yv[i];
				for (int j = Ap[i]; j < Ap[i + 1]; ++j)
					temp -= Av[j] * xv[Ai[j]];
				wv[i] = temp;
			}
		}
	}
	else
	{
		beta /= alpha;

		if (beta == 0)
		{
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
			for (int i = 0; i < n; ++i)
			{
				double temp = 0.0;
				for (int j = Ap[i]; j < Ap[i + 1]; ++j)
					temp += Av[j] * xv[Ai[j]];
				wv[i] = alpha * temp;
			}
		}
		else if (beta == 1)
		{
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
			for (int i = 0; i < n; ++i)
			{
				double temp = yv[i];
				for (int j = Ap[i]; j < Ap[i + 1]; ++j)
					temp += Av[j] * xv[Ai[j]];
				wv[i] = alpha * temp;
			}
		}
		else if (beta == -1)
		{
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
			for (int i = 0; i < n; ++i)
			{
				double temp = -yv[i];
				for (int j = Ap[i]; j < Ap[i + 1]; ++j)
					temp += Av[j] * xv[Ai[j]];
				wv[i] = alpha * temp;
			}
		}
		else
		{
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
			for (int i = 0; i < n; ++i)
			{
				double temp = beta * yv[i];
				for (int j = Ap[i]; j < Ap[i + 1]; ++j)
					temp += Av[j] * xv[Ai[j]];
				wv[i] = alpha * temp;
			}
		}
	}
}

void CSR_MatVec(const CSR_Matrix& A, const Vector& x, const Vector& y)
{
	CSR_MatVec(1.0, A, x, 0.0, y, y);
}

}